package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"github.com/gin-gonic/gin"
)

const (
    BrokerAddress    = "localhost"
    BrokerPort       = 1883
    MqttPublishTopic = "response"
)

type RequestData struct {
	Data struct {
		UserID string `json:"user_id"`
	} `json:"data"`
	UUID string `json:"uuid"`
}

type ResponseData struct {
	UserID   string `json:"user_id"`
	Response int    `json:"response"`
}

type CANRequest struct {
	CANID     uint32    `json:"can_id"`
	VehicleID string    `json:"vehicle_id"`
	Data      []interface{}    `json:"data"`
	Timestamp int64     `json:"timestamp"`
}

type CANResponse struct {
	CANID      uint32  `json:"can_id"`
	VehicleID  string  `json:"vehicle_id"`
	AttackType int     `json:"attack_type"`
	AttackName string  `json:"attack_name"`
	Confidence float64 `json:"confidence"`
	Timestamp  int64   `json:"timestamp"`
}

var (
	mqttClient      mqtt.Client
	mqttInitialized = make(chan struct{})
)

func onConnect(client mqtt.Client) {
	log.Println("Connected to MQTT broker")
	close(mqttInitialized)
}

func setupMQTTClient() error {
	opts := mqtt.NewClientOptions()
	brokerURL := fmt.Sprintf("tcp://%s:%d", BrokerAddress, BrokerPort)
	opts.AddBroker(brokerURL)
	opts.SetClientID("Worker2")
	opts.SetCleanSession(true)
	opts.SetKeepAlive(60 * time.Second)
	opts.SetAutoReconnect(true)
	opts.SetMaxReconnectInterval(60 * time.Second)

	opts.SetOnConnectHandler(onConnect)
	opts.SetConnectionLostHandler(func(client mqtt.Client, err error) {
		log.Printf("Connection lost: %v", err)
	})

	mqttClient = mqtt.NewClient(opts)
	if token := mqttClient.Connect(); token.Wait() && token.Error() != nil {
		return fmt.Errorf("Failed to connect to MQTT broker: %v", token.Error())
	}
	
	<-mqttInitialized
	
	return nil
}

func handleRequest(c *gin.Context) {
	var requestData RequestData
	if err := c.ShouldBindJSON(&requestData); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid JSON"})
		return
	}

	fmt.Printf("Received data for user ID: %s\n", requestData.Data.UserID)

	time.Sleep(2 * time.Second)

	response := ResponseData{
		UserID:   requestData.UUID,
		Response: 200,
	}

	respJson, err := json.Marshal(response)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to serialize response"})
		return
	}

	topic := fmt.Sprintf("%s/%s", MqttPublishTopic, requestData.UUID)
	token := mqttClient.Publish(topic, 0, false, respJson)
	if token.Wait() && token.Error() != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to publish to MQTT"})
		return
	}

	c.JSON(http.StatusOK, response)
}

func handleCANRequest(c *gin.Context) {
	var canRequest CANRequest
	if err := c.ShouldBindJSON(&canRequest); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid CAN request format"})
		return
	}

	log.Printf("Received CAN message: ID=%d, Vehicle=%s", canRequest.CANID, canRequest.VehicleID)

	time.Sleep(2 * time.Second)

	response := CANResponse{
		CANID:      canRequest.CANID,
		VehicleID:  canRequest.VehicleID,
		AttackType: 2,
		AttackName: "Brute Force",
		Confidence: 0.87,
		Timestamp:  time.Now().Unix(),
	}

	respJson, err := json.Marshal(response)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to serialize response"})
		return
	}

	topic := fmt.Sprintf("%s/vehicle:%s:can:%d:%d", MqttPublishTopic, 
		canRequest.VehicleID, canRequest.CANID, canRequest.Timestamp)
	
	token := mqttClient.Publish(topic, 0, false, respJson)
	if token.Wait() && token.Error() != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to publish to MQTT"})
		return
	}

	c.JSON(http.StatusOK, response)
}

func main() {
	if len(os.Args) < 2 {
		log.Fatal("Port number must be provided as the first argument")
	}

	port, err := strconv.Atoi(os.Args[1])
	if err != nil {
		log.Fatalf("Invalid port number: %v", err)
	}

	mqttInitialized = make(chan struct{})
	
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		if err := setupMQTTClient(); err != nil {
			log.Fatalf("Failed to set up MQTT: %v", err)
		}
	}()

	router := gin.Default()
	router.POST("/simulate", handleRequest)
	router.POST("/can/analyze", handleCANRequest)

	wg.Wait()

	serverAddr := fmt.Sprintf("0.0.0.0:%d", port)
	log.Printf("Starting server on %s", serverAddr)
	if err := router.Run(serverAddr); err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}
